import time
from threading import Thread
from robomaster import robot
from robomaster import camera
import websocket
import _thread
import cv2
import demjson
import asyncio
import base64


ep_robot = robot.Robot()
ep_robot.initialize(conn_type="sta")

ep_camera = ep_robot.camera


def on_open(ws):
    ws.send('{"type":"login","client_name":"sender","id":"2"}')
    _thread.start_new_thread(video_stream_send())
    # _thread.start_new_thread(ws.run_forever())



def on_message(wsapp, message):
    data = demjson.decode(message)
    if data["type"] == 'control':
        x_val = 0.5
        y_val = 0.3
        z_val = 30
        #{"type":"control","direction":""}
        ep_chassis = ep_robot.chassis
        print(message)
        if data["direction"] == 'forward':
            ep_chassis.drive_speed(x=x_val, y=0, z=0, timeout=5)
        elif data["direction"] == 'backward':
            ep_chassis.drive_speed(x=-x_val, y=0, z=0, timeout=5)
        elif data["direction"] == 'to_left':
            ep_chassis.drive_speed(x=0, y=-y_val, z=0, timeout=5)
        elif data["direction"] == 'to_right':
            ep_chassis.drive_speed(x=0, y=y_val, z=0, timeout=5)
        elif data["direction"] == 'turn_left':
            ep_chassis.drive_speed(x=0, y=0, z=-z_val, timeout=5)
        elif data["direction"] == 'turn_right':
            ep_chassis.drive_speed(x=0, y=0, z=z_val, timeout=5)
        elif data["direction"] == 'stop':
            ep_chassis.drive_speed(x=0, y=0, z=0, timeout=5)

    else:
        print(message)

def on_ping(ws):
    ws.send('{"type":"test","client_name":"sender","id":"2"}')

def on_close(ws):
    cv2.destroyAllWindows()
    ep_camera.stop_video_stream()
    ep_robot.close()


def image_to_base64(image_np):
    image = cv2.imencode('.jpg', image_np)[1]
    image_code = str(base64.b64encode(image))[2:-1]
    return image_code


def video_stream_send():
    ep_camera.start_video_stream(display=False)

    while True:

        img = ep_camera.read_cv2_image()
        img = cv2.resize(img,(500,300))
        encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 15]
        base64_str = cv2.imencode('.jpg', img, encode_param)[1].tobytes()
        # s = base64.b64encode(base64_str)[2:-1]
        s = image_to_base64(img)
        ws.send('{"type":"stream","client_name":"sender","id":"2","base64_image":"%s"}' %s)
        # ws.send('{"type":"stream","client_name":"sender","id":"2","base64_image":"123"}')
        # print('send image success')
        time.sleep(0.1)
        cv2.imshow("Robot", img)
        cv2.waitKey(1)


# ws = websocket.WebSocketApp("ws://zhikuxf.com:7272", on_message=on_message, on_open=on_open, on_ping=on_ping, on_close=on_close)
ws = websocket.WebSocketApp("ws://127.0.0.1:7272", on_message=on_message, on_open=on_open, on_ping=on_ping, on_close=on_close)
_thread.start_new_thread(ws.run_forever())


